# 读取rosbag文件，将路径保存到本地

import argparse
import cv2
import os
import numpy as np

import rospy
import rosbag
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped


def save_path_as_kitti(path_msg, output_file):
    pass

def save_path_as_kitti_(path_msg, output_file):
    pass


def save_path_as_tum(path_msg, output_file):
    with open(output_file, 'w') as file:
        for pose in path_msg.poses:
            # TUM格式：timestamp x y z qx qy qz qw
            timestamp = pose.header.stamp.to_sec()
            x = pose.pose.position.x
            y = pose.pose.position.y
            z = pose.pose.position.z
            qx = pose.pose.orientation.x
            qy = pose.pose.orientation.y
            qz = pose.pose.orientation.z
            qw = pose.pose.orientation.w
            file.write(f"{timestamp} {x} {y} {z} {qx} {qy} {qz} {qw}\n")


def save_path_as_tum_(msg, file):
    # TUM格式：timestamp x y z qx qy qz qw
    timestamp = msg.header.stamp.to_sec()
    x = msg.pose.position.x
    y = msg.pose.position.y
    z = msg.pose.position.z
    qx = msg.pose.orientation.x
    qy = msg.pose.orientation.y
    qz = msg.pose.orientation.z
    qw = msg.pose.orientation.w
    file.write(f"{timestamp} {x} {y} {z} {qx} {qy} {qz} {qw}\n")


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', type=str,
                        default="/home/daybeha/Documents/github/Hybird-NeuroSLAM/vins_ws/VIO_urban39_keyframe.bag")
                        # default="/home/daybeha/Documents/github/Hybird-NeuroSLAM/ORB_SLAM3_detailed_comments-master/Examples/src/stereo_kitti08.bag")
                        # default="/media/daybeha/LTFM2/dataset/Event_cam/Vector/small/with_depth/sofa_normal1.bag")
    parser.add_argument('--topics', nargs='+', default=['/vins/path'])
    # parser.add_argument('--topics', nargs='+', default=['/orb/path'])
    # parser.add_argument('--topics', nargs='+', default=['/gt/pose'])
    parser.add_argument('--save_dir', type=str,
                        default="/home/daybeha/Documents/github/Hybird-NeuroSLAM/vins_ws/results/")
                        # default= "/home/daybeha/Documents/github/Hybird-NeuroSLAM/ORB_SLAM3_detailed_comments-master/Examples/src/results/")
                        # default="/media/daybeha/LTFM2/dataset/Event_cam/Vector/small/with_depth/gt/")

    parser.add_argument('--traj_format', type=str, default="tum", help='kitti or tum')
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = get_args()
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)

    # 打开ROS bag文件
    bag = rosbag.Bag(args.bag)

    # 获取不带后缀的文件名
    bag_name = os.path.splitext(os.path.basename(args.bag))[0]
    # file_save = os.path.join(args.save_dir, bag_name + "_gt.txt")
    file_save = os.path.join(args.save_dir, bag_name + ".txt")

    timestamp_list = []
    # save_dir = args.save_dir
    traj_format = args.traj_format

    end_time = bag.get_end_time()
    msg = bag.read_messages(topics=args.topics, start_time=rospy.Time.from_sec(end_time - 2))
    last_msg = None

    # 从bag末尾开始倒序查找最近的'/orb/path'消息
    for topic, msg, t in msg:
        last_msg = msg


    target_type = bag.get_type_and_topic_info().topics[args.topics[0]].msg_type
    if target_type == "nav_msgs/Path":
        if last_msg:
            if traj_format == "kitti":
                save_path_as_kitti(last_msg, file_save)
            elif traj_format == "tum":
                save_path_as_tum(last_msg, file_save)
            print(f"Saved last path to {file_save}")
    elif target_type == "geometry_msgs/PoseStamped":
        if traj_format == "kitti":
            with open(file_save, 'w') as file:
                for topic, msg, t in bag.read_messages(topics=args.topics):
                    save_path_as_kitti_(msg.pose, file)
        elif traj_format == "tum":
            with open(file_save, 'w') as file:
                for topic, msg, t in bag.read_messages(topics=args.topics):
                    save_path_as_tum_(msg, file)
        print(f"Saved path to {file_save}")

    # 关闭ROS bag文件
    bag.close()
